import re
from scapy.all import rdpcap
import pandas as pd
import scapy
import binascii
import pandas as pd
import scapy.all as scapy
from functools import reduce
from flowcontainer.extractor import extract
import os
from scapy.all import rdpcap
"""
    CSV Utils
"""
class CSVUtils :
    
    """ 
        find ip from params
    """
    @staticmethod
    def find_ip(fun:str,params:str) -> tuple[list,bool]:
        results = []
        """  
        match connect() parameters
            ex: ...htons(443),....::ffff:40.70.161.102...
            return port and ip
            ipv4 or ipv6
        """
        connect_ipv4_pattern = r'htons\((\d+)\).*?::ffff:(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})'
        connect_ipv6_pattern = r'htons\((\d+)\).*?inet_pton\(AF_INET6,"([\w:]+)"'

        """
        match connected socket
            ex:63<TCPv6:[[::ffff:10.0.2.16]:39830->[::ffff:104.18.215.59]:443]>
            return source ip\port  des ip\port
        """
        compl_ipv4_pattern = r'::ffff:(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\]:(\d+).*?::ffff:(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\]:(\d+)'
        """ 
        match ipv6 
            ex:<UDPv6:[[fec0::41dd:35fd:9649:f8cf]:51994->[2001:4860:4860::8888]:53]>
        """
        compl_ipv6_pattern = r'\[(\w{0,4}(:\w{0,4}){1,7})\]:(\d+)->\[(\w{0,4}(:\w{0,4}){1,7})\]:(\d+)'

        """  
        match content of sendmsg \ revfrom 
            but it is uncompleted
            return portion of a load
        """
        if fun == 'connect':
            match_connect_ipv4_pattern = re.search(connect_ipv4_pattern,params)
            if match_connect_ipv4_pattern:
                port = match_connect_ipv4_pattern.group(1)
                ip = match_connect_ipv4_pattern.group(2)
                results.append((port,ip))
                return results,True
            match_connect_ipv6_pattern = re.search(connect_ipv6_pattern,params)
            if match_connect_ipv6_pattern:
                port = match_connect_ipv6_pattern.group(1)
                ip = match_connect_ipv6_pattern.group(2)
                results.append((port,ip))
                return results,True

        elif fun == 'sendmsg' or fun == 'sendto' or fun == 'recvfrom' or fun == 'recvmsg':

            match_compl_ipv4 = re.search(compl_ipv4_pattern, params)
            if match_compl_ipv4:
                src_ip = match_compl_ipv4.group(1)
                src_port = match_compl_ipv4.group(2)
                dest_ip = match_compl_ipv4.group(3)
                dest_port = match_compl_ipv4.group(4)
                results.append((src_ip, src_port, dest_ip, dest_port))
                return results,True
            match_compl_ipv6 = re.search(compl_ipv6_pattern, params)
            if match_compl_ipv6:
                src_ip = match_compl_ipv6.group(1)
                src_port = match_compl_ipv6.group(3)
                dest_ip = match_compl_ipv6.group(4)
                dest_port = match_compl_ipv6.group(6)
                results.append((src_ip, src_port, dest_ip, dest_port))
                return results,True
            
            """  
                Beacuse these system calls maybe have linke .htons(443),....::ffff:40.70.161.102...
            """
            match_connect_ipv4_pattern = re.search(connect_ipv4_pattern,params)
            if match_connect_ipv4_pattern:
                port = match_connect_ipv4_pattern.group(1)
                ip = match_connect_ipv4_pattern.group(2)
                results.append((port,ip))
                return results,True
            match_connect_ipv6_pattern = re.search(connect_ipv6_pattern,params)
            if match_connect_ipv6_pattern:
                port = match_connect_ipv6_pattern.group(1)
                ip = match_connect_ipv6_pattern.group(2)
                results.append((port,ip))
                return results,True
        return '',False
    
    """  
        find sendmsg,sendto,recvform,recvmsg about info of traffic 
    """

    @staticmethod
    def find_params_content(params:str) -> str :
        content_pattern  = r'"((?:[^"\\]|\\.)*?)"'
        content_match = re.search(content_pattern,params)
        if content_match :
            return content_match.group(1)
    
        
""" 
    Traffic Utils
"""
class TrafficUtils :
    def __init__(self, word_dir=None,word_name=None):
        self.word_dir = word_dir
        self.word_name = word_name

    """
        from ET-BERT's code  
            extract burst of traffic future
    """

    @staticmethod
    def split_cap(father_path,pcap_name, dataset_level='flow'):
        pcap_file = os.path.join(father_path,pcap_name)
        # 根据dataset_level选择命令
        if dataset_level == 'flow':
            cmd = f"mono /home/chenjiaxiong/jx_code/syscall-traffic/SplitCap.exe -r {pcap_file} -s session -o {father_path}"
        elif dataset_level == 'packet':
            cmd = f"mono /home/chenjiaxiong/jx_code/syscall-traffic/SplitCap.exe -r {pcap_file} -s packets 1 -o {father_path}"

        # 执行命令
        os.system(cmd)
        os.remove(pcap_file)
        return father_path
    
    @staticmethod
    def convert_to_compressed_log_format(byte_data):
        result = []
        escape_dict = {
        7: r'\a', 8: r'\b', 9: r'\t',
        10: r'\n', 11: r'\v', 12: r'\f', 13: r'\r', 32: r' '}

        for byte in byte_data:
            # 如果是可打印的ASCII字符（0x20到0x7E），直接转换为字符
            char = chr(byte)
            if char == '"':
                # 将双引号 " 转义为 \"
                result.append('\\"')
            elif char == '\\':
                # 将反斜杠 \ 转义为 \\
                result.append('\\\\')
            elif byte in escape_dict:
                result.append(escape_dict[byte])
            elif 32 <= byte <= 126:
                # 对于其他可打印字符，直接添加
                result.append(char)
            else:
                # 不可打印字符，转换为压缩的八进制表示，去除前导零
                result.append(f"\\{int(oct(byte)[2:])}")
        return ''.join(result)

        

    def cut(self,obj, sec):
        result = [obj[i:i+sec] for i in range(0,len(obj),sec)]
        try:
            remanent_count = len(result[0])%4
        except Exception as e:
            remanent_count = 0
            print("cut datagram error!")
        if remanent_count == 0:
            pass
        else:
            result = [obj[i:i+sec+remanent_count] for i in range(0,len(obj),sec+remanent_count)]
        return result

    def bigram_generation(self,packet_datagram, packet_len = 64, flag=True):
        result = ''
        generated_datagram = self.cut(packet_datagram,1)
        token_count = 0
        for sub_string_index in range(len(generated_datagram)):
            if sub_string_index != (len(generated_datagram) - 1):
                token_count += 1
                if token_count > packet_len:
                    break
                else:
                    merge_word_bigram = generated_datagram[sub_string_index] + generated_datagram[sub_string_index + 1]
            else:
                break
            result += merge_word_bigram
            result += ' '
        
        return result

    def get_burst_feature(self,label_pcap, payload_len=64):
        packets = rdpcap(label_pcap)
        packet_direction = []
        feature_result = extract(label_pcap,filter="(tcp or udp or gre) or tcp.analysis.retransmission" )
        for key in feature_result.keys():
            value = feature_result[key]
            packet_direction = [x // abs(x) for x in value.ip_lengths]
        payloads_per_burst = []
        if len(packet_direction) == len(packets):
            burst_data_string = ''
            burst_txt = ''

            for packet_index, packet in enumerate(packets):
                """  
                    这里要注意：这里可能可以直接提取payload字段！哥们。注意
                """

                   # 提取TCP/UDP的payload
                if packet.haslayer('TCP') or packet.haslayer('UDP'):
                    if packet.haslayer('Raw'):
                        payload = packet['Raw'].load
                        truncated_payload = payload[:128] if len(
                            payload) >= 128 else payload
                        # 将payload加入当前 burst_feature 对应的 payload 列表
                        result = self.convert_to_compressed_log_format(truncated_payload)
                        payloads_per_burst.append(result)
                packet_data = packets[packet_index].copy()
                data = (binascii.hexlify(bytes(packet_data)))
                
                packet_string = data.decode()[:2*payload_len]

                if packet_index == 0:
                    burst_data_string += packet_string

                else:
                    if packet_direction[packet_index] != packet_direction[packet_index - 1]:
                        
                        length = len(burst_data_string)
                        if len(payloads_per_burst)!= 0:
                            burst_txt += '\n'.join([str(payload) for payload in payloads_per_burst]) + '\n'
                        else :
                            continue
                        payloads_per_burst = []
                        for string_txt in self.cut(burst_data_string, int(length / 2)):
                            burst_txt += self.bigram_generation(string_txt, packet_len=len(string_txt))
                            burst_txt += '\n'
                        burst_txt += '\n'
                        
                        burst_data_string = ''
                    
                    burst_data_string += packet_string
                    if packet_index == len(packets) - 1:
                        length = len(burst_data_string)
                        if len(payloads_per_burst)!= 0:
                            burst_txt += '\n'.join([str(payload) for payload in payloads_per_burst]) + '\n'
                        else:
                            continue
                        payloads_per_burst = []
                        for string_txt in self.cut(burst_data_string, int(length / 2)):
                            burst_txt += self.bigram_generation(string_txt, packet_len=len(string_txt))
                            burst_txt += '\n'
                        burst_txt += '\n'
            file_path = os.path.join(self.word_dir, self.word_name)
            with open(file_path,'a') as f:
                f.write(burst_txt)
        return 0
    
    @staticmethod
    def is_hex_format(line):
        # 定义匹配至少两个以上的四字符16进制块的正则表达式
        hex_pattern = re.compile(r'^([0-9a-fA-F]{4}\s+){2,}[0-9a-fA-F]{4}$')

        # 检测是否匹配
        if hex_pattern.match(line.strip()):
            return True
        else:
            return False